import json
import traceback
import warnings
import os

from elasticsearch import Elasticsearch
from elasticsearch import helpers
from elasticsearch.exceptions import ElasticsearchWarning

# Suppress Elasticsearch security warnings
warnings.filterwarnings('ignore', category=ElasticsearchWarning)

law_data = [
    "第一条 为切实加强数据安全管理，确保数据依法合规有效使用，进一步提升合作方数据安全风险防范和保障能力，依据《中华人民共和国数据安全法》《中华人民共和国个人信息保护法》等法律法规，以及《中国联通数据安全管理办法》（中国联通集团〔2024〕50 号）《中国联通个人信息保护管理办法（试行）》（中国联通集团〔2022〕118号）等公司规定，制定本办法。",
    "第二条 本办法适用于中国联合网络通信集团有限公司及其分公司和全资、控股子公司，中国联合网络通信有限公司及其分公司和全资、控股子公司（以下简称“各单位”）。",
    "第三条 本办法所称合作方是指参与业务合作、技术支撑、数据服务等可能接触到各单位所管理和控制的数据的外部单位。其中，业务合作主要包括数据业务合作推广、渠道接入等形式；技术支撑主要包括系统开发集成、系统维护、系统测试等形式；数据服务指数据全生命周期处理过程中所提供的服务，主要包括数据提取、数据建模、数据挖掘、数据分析等形式。",
    "第四条 数据按照业务属性分为用户相关数据和企业自身数据两大类；根据数据遭到篡改、破坏、泄露或者非法获取、非法利用，对国家安全、公共利益或者个人、组织合法权益等造成的危害程度，将数据分为一般数据、重要数据和核心数据。具体数据分类分级规则以集团公司数据分类分级管理规定为准。",
    "第五条 坚持敏感数据“不出门”，除法律法规规定或获得个人信息主体明确授权外，不允许向合作方提供敏感数据。严格遵循合法、正当、必要和诚信的原则，在个人信息主体授权同意的范围内与合作方开展合作，严禁合作方超出合同约定的目的和范围处理用户个人信息，或使用匿名化处理之后的数据。敏感数据的界定以集团公司数据分类分级管理规定为准。",
    "第六条 重要数据和核心数据原则上不得对合作方提供，如因特殊情况确需提供的，按照《中国联通数据安全管理办法》的相关规定执行。",
    "第七条 集团公司网络与信息安全部是公司数据安全工作的统筹管理单位，负责制定并更新合作方数据安全管理制度，对各单位的合作方数据安全管理情况和落实效果进行监督检查。",
    "第八条 各单位网络与信息安全部门为本单位数据安全工作的统筹管理部门，具体职责包括：\n（一）对本单位合作方的数据安全保护能力进行审核；\n（二）报备本单位数据业务相关合作协议至总部数据安全运营系统；\n（三）汇总并动态更新本单位合作方清单；\n（四）及时上报本单位合作方数据安全事件至集团公司网络与信息安全部；\n（五）组织开展数据合作安全管理培训，并对培训情况进行记录。",
    "第九条 按照“谁引入谁负责”及“谁管业务、谁管业务数据、谁管数据安全”的原则，各单位合作方引入部门是本单位合作方数据安全管理责任单位，具体职责包括：\n（一）对合作方及合作内容进行合规性评估；\n（二）与合作方签署合作协议、承诺书等法律文件；\n（三）梳理形成并定期更新本部门涉及的合作方清单，及时向本单位数据安全统筹管理部门报备；\n（四）发起合作方系统与本单位系统的对接审批流程，明确数据内容和合作方人员情况；\n（五）对合作方数据使用情况进行审计、行为约束和监督管理；\n（六）建立本部门应急响应预案，完善应急响应机制，及时处置合作方数据安全事件。",
    "第十条 各单位向合作方提供数据的部门为数据输出部门，对提供的数据承担安全管理责任，具体职责包括：\n（一）梳理形成并定期更新外部系统对接台账，定期向本单位数据安全统筹管理部门报备；\n（二）采用认证鉴权、访问控制等技术手段，严格管理合作方账号、合作方终端、合作方调用接口；\n（三）强化数据防泄漏、数据脱敏、数据安全风险监测等技术措施，定期开展巡查巡检；\n（四）建立合作方数据访问、使用和传输行为审计机制，并留存审计记录。",
    "第十一条 合作开展前，合作方引入部门应对合作方运营资质、人员能力等进行资质审核，对合作方数据来源合法性、数据安全事件情况、违反数据安全法律法规情况及数据合作信誉情况等开展调查。合作方资质审核与背景调查事项清单详见附件。",
    "第十二条 涉及合作方处理个人信息或重要数据、核心数据的，在合作开展前，应由合作方提供数据安全风险评估报告（自评估或第三方评估），合作方引入部门应对合作方的数据安全保护能力进行核验，并将核验结果提交至本单位数据安全统筹管理部门进行审核，审核通过后方可认为其具备合作资质。涉及合作方处理核心数据的，合作方引入部门应当评估安全风险，采取必要的安全保护措施，并由属地行业监管部门审查后报工业和信息化部，工业和信息化部按照有关规定进行审查，审查通过后方可处理。",
    "第十三条 针对参与本单位重要数据、核心数据处理活动的合作方人员，合作方引入部门应对合作方人员进行背景调查，包括但不限于身份信息验证、教育背景验证、工作履历核实、犯罪记录查询、法院失信记录查询、职业资格证书验证等。若提供信息不准确、验证不通过或存在犯罪记录、法院失信记录的，原则上合作方人员不应参与本单位重要数据、核心数据处理活动。",
    "第十四条 各单位应与其引入的合作方签订合作协议，落实合作方数据安全保护责任，协议中应当明确如下内容：\n（一）数据处理活动的场景、数据类型、处理目的、处理期限、处理方式、合作方名称及联系方式；\n（二）合作方对其提供数据来源及其数据处理活动的合法合规性承担保证责任；\n（三）按照最小必要原则、结合数据合作等情况，允许合作方人员接触到的数据处理相关平台系统，及其数据使用的权限、内容、范围及用途；\n（四）合作方数据安全保护责任，包括但不限于合作方对接触及获取的数据履行保密义务；不得未经同意改变数据使用的目的、方式和范围；不得缓存、窃取、泄露、滥用、非法向他人提供数据；不得对数据进行加工后还原成原始数据；合作方的数据安全保护措施不得低于本单位的数据安全保护要求；应遵守本单位相关管理办法，并接受必要的安全监控和审计等；\n（五）针对合作事项涉及的数据安全风险，合作方应配备的安全保护措施，如数据加密传输、数据加密存储、敏感数据脱敏、操作权限控制、日志留存、操作行为审计等；\n（六）合作结束后，合作方应及时返还相关资料、销毁数据；\n（七）合作方违反合同约定时应承担的相关责任，责任内容应具体、清晰、可执行；\n（八）合作方对于可接触用户个人信息或重要数据、核心数据的员工的管理要求，包括合作方人员的数据安全保护责任、合作方人员发生变动应及时告知并应签订个人安全保密承诺书等，合作方人员的安全责任不得低于合作协议规定程度；\n（九）是否允许合作方向境外传输、存储合作中涉及的数据，以及数据跨境相关限制性要求；\n（十）涉及个人信息或重要数据、核心数据的，合作方应每年就合作内容至少开展一次数据安全风险评估并提交评估报告；\n（十一）依据法律法规、公司制度规定及具体业务情况需要约定的其他事项。",
    "第十五条 各单位通过间接途径获取数据的，包括在数据对外合作过程中，合作方采集、提供外部数据等情况，合作方引入部门应要求合作方应说明数据来源，并签署相关书面承诺，确保外部数据采集渠道的合法合规且不存在侵权事项，同时还应出示个人授权协议，确保数据使用未超出授权范围。",
    "第十六条 各单位应在签署合作协议后 10 个工作日内，在总部数据安全运营系统进行备案，备案内容包括业务名称、业务介绍、合作对象、合作数据范围、是否涉及敏感数据输出、是否进行业务数据安全风险评估、技术能力保障情况等，不包括数据内容本身。",
    "第十七条 各单位应汇总形成并动态更新本单位合作方清单，包括合作方单位名称、联系人信息，以及合作形式、期限、合作业务或系统、合作数据范围、输出数据类型、数据输出方式、数据输出部门、数据安全保护措施等信息。",
    "第十八条 合作方系统需接入本单位系统时，应满足《中国联通合作方网络与信息安全管理办法》（网信安〔2023〕22 号）及本单位网络与信息安全管理要求。由合作方引入部门负责发起合作方系统与本单位系统的对接流程，由数据输出部门将系统接入情况提交本单位数据安全统筹管理部门备案，并保留系统对接台账和系统日志至少 6个月。",
    "第十九条 各单位应对接入本单位内部网络的合作方终端进行严格接入认证。接入内部网络的设备应限制对网口和无线上网的使用，并安装符合要求的防病毒软件。",
    "第二十条 数据输出部门应对合作方接口调用的情况进行严格审核，严禁合作方超范围进行接口调用。合作方系统退出服务时，数据输出部门应及时关闭相应接口。",
    "第二十一条 数据输出部门应对涉及核心数据、重要数据、用户个人信息处理活动的合作方账号实行严格管理：\n（一）原则上，禁止合作方人员掌握数据提取相关系统管理员权限，禁止为合作方人员分配具备创建系统账号或者其他超出工作范围权限的高权限账号；合作方账号应实人实名，账号不得转让、转借；\n（二）每季度对全量合作方账号进行排查审核，控制合作方涉及敏感数据的人员范围；未经审批，禁止对合作方等非本单位人员分配涉及敏感数据的权限账号；\n（三）按照《中国联通信息系统账号运营安全管理规范》（数字化〔2022〕16 号）有关规定，严格管理合作方账号申请、回收、授权等环节，确保合作方人员发生离职或岗位变动时及时清理其账号，数据合作结束后及时回收账号权限。",
    "第二十二条 数据输出部门应对合作方人员采用严格访问控制策略和管控手段，经合作方引入部门和数据输出部门审批通过后，合作方人员方可进入可能接触到个人信息或重要数据、核心数据的生产或维护区域。",
    "第二十三条 合作方人员转岗或离岗前，合作方引入部门应根据相关规定要求合作方提交人员转岗或离岗说明，完成合作方人员的账号回收、审核、网络调整等工作后，合作方人员方可转岗或离岗。",
    "第二十四条 同一业务涉及多个合作方时，合作方引入部门应保证每个合作方只获得其工作所需的数据资源。未经合作方引入部门同意，合作方之间不得进行数据传输、交换。",
    "第二十五条 数据输出部门应强化数据防篡改、防泄漏以及数据脱敏等安全技术措施，针对服务器、数据库等每月开展巡查巡检。",
    "第二十六条 数据输出部门应建立合作方数据访问、使用和传输行为审计机制，基于操作权限、接口权限管控规则对异常数据访问、使用、传输等进行人工审计或自动审计，并留存审计记录。",
    "第二十七条 数据输出部门应建立数据安全风险监测机制，实现对合作方数据访问行为和使用行为的监测、分析与异常预警，及时发现可能存在的数据安全风险。",
    "第二十八条 各单位面向合作方开放的数据接口应具备接口认证鉴权与安全监控能力，能够限制违规设备接入，对接口调用进行必要的自动监控和处理，监控可能存在的数据异常调用、数据违规获取等风险。",
    "第二十九条 合作期限内，合作方引入部门应要求合作方对数据访问行为及数据传输进行日志留存，操作日志至少包括数据操作人员、操作数据类型、数据操作行为、数据操作时间等；接口日志至少包括发起者和数据接收者、接口名称、调用时间、调用目的、数据名称、数据量级等，日志留存确保满足审计和溯源要求。合作方引入部门应要求合作方至少每季度报备合作数据实际使用情况，报备内容至少包括数据使用范围、数据使用人员、数据使用方式、数据使用目的、数据使用日期（区间）。合作方引入部门应对合作方数据使用情况报备内容进行审核，确保符合法律法规和合同约定。",
    "第三十条 合作方引入部门应对合作方数据使用情况进行审计、行为约束和监督管理。",
    "第三十一条 按照法律法规、政府部门要求等基于特定场景需要提供数据的，合作方引入部门应落实公司保密规定、数据安全管理要求，并保存相关授权/证明文件、数据输出台账以作备案。基于特定场景的数据提取，数据输出部门应建立专门账号集中管理，保证账号专人专用，保留此账号的所有操作日志便于后续涉敏审计。",
    "第三十二条 原则上，各单位不得允许合作方向境外传输、存储数据，如基于合作事项确需向境外传输、存储数据的，应经合作方引入单位同意，合作方应按照《数据出境安全评估办法》（国家互联网信息办公室令第 11 号）等国家法律法规规定，开展数据出境风险自评估和安全评估，采取必要的安全措施，在确保安全的前提下进行数据出境，并加强对数据出境后的跟踪掌握；核心数据不得出境。",
    "第三十三条 合作方引入部门应在合作结束后，采取有效措施督促合作方按照合同约定返还资料、销毁数据，不得超期留存。由合作方提供销毁记录承诺书以作备案，承诺书中应明确数据未按要求销毁的法律责任。",
    "第三十四条 合作方引入部门应建立应急响应预案，完善应急响应机制，如出现合作方数据安全事件，应第一时间进行有效处置，做好记录，并上报本单位数据安全统筹管理部门，数据安全统筹管理部门应及时上报集团公司网络与信息安全部。",
    "第三十五条 各单位应每年开展数据合作安全管理培训，培训内容包括但不限于数据安全相关法律法规、公司数据安全制度、保密责任及义务、安全建设与运营规范、安全意识等，并对培训情况进行记录。",
    "第三十六条 集团公司网络与信息安全部每年采取自查和抽查的方式，对省级分（子）公司数据安全管理情况和落实效果进行监督检查，对检查结果进行全国通报，并及时督促问题整改。",
    "第三十七条 合作方引入部门应加强合作方数据安全监督管理，包括但不限于对数据合作清单进行定期核查、及时更新，确保与数据合作实际情况保持一致；对合作方数据安全保护措施进行动态监管，要求合作方提供数据使用用途或安全审计记录等，及时发现存在的安全风险并督促问题整改。",
    "第三十八条 合作方未履行合作协议约定，导致数据安全事件发生，造成不良影响的，经相关引入单位网络安全领导小组审议通过后，将相关企业纳入合作方数据安全黑名单和供应商黑名单管理。根据《中国联通供应商黑名单管理办法（修订）》（中国联通集团〔2021〕150号），由集团公司物资采购与管理部落实数据合作项目采购时的供应商筛选要求，其他合作形式的合作方筛选由各单位自行落实。合作方数据安全黑名单应及时上报集团公司网络与信息安全部，列入黑名单的企业，各单位不得与其开展涉及数据合作的业务。如遇特殊情况，报集团公司网络与信息安全部一事一议。",
    "第三十九条 合作方或其员工违反国家法律法规、公司相关规章制度或合作协议约定，导致数据安全事件发生，造成不良影响的，有关单位应依据合同约定进行追责。涉嫌违法犯罪的，移送司法机关依法处理。",
    "第四十条 各单位自有人员因数据安全工作存在落实不到位、疏于管理、失职等违规行为导致数据安全事件发生，造成不良影响的，有关单位应依据《中国联通网络与信息安全系统安全工作奖惩办法》进行严肃问责。涉嫌违法犯罪的，移送司法机关依法处理。"
]

# 创建 Elasticsearch 客户端实例
# 从环境变量读取连接信息

ES_HOST = os.getenv('ES_HOST', '1.94.252.115')
ES_PORT = int(os.getenv('ES_PORT', '9200'))
ES_USERNAME = os.getenv('ES_USERNAME', 'admin')
ES_PASSWORD = os.getenv('ES_PASSWORD', '383448732e8ea746216d')
ES_SCHEME = os.getenv('ES_SCHEME', 'http')

# 根据是否配置了认证信息来决定连接方式
print(f"ES连接配置: {ES_SCHEME}://{ES_HOST}:{ES_PORT}")
print(f"ES认证信息: username={ES_USERNAME}, password={'***' if ES_PASSWORD else None}")

# ES 8.x版本需要特殊配置 - 使用URL格式包含认证信息
es = Elasticsearch(
    f"{ES_SCHEME}://{ES_USERNAME}:{ES_PASSWORD}@{ES_HOST}:{ES_PORT}",
    verify_certs=False,
    ssl_show_warn=False,
    request_timeout=30
)
print("ES客户端创建成功(URL认证方式)")


def insert_data(index_name: str, data: list[dict]):
    global es
    try:
        print("正在尝试连接ES...")
        # 检查连接状态
        if not es.ping():
            print('无法连接到 Elasticsearch，尝试重新连接...')
            # 重新创建连接
            es = Elasticsearch(
                f"{ES_SCHEME}://{ES_USERNAME}:{ES_PASSWORD}@{ES_HOST}:{ES_PORT}",
                verify_certs=False,
                ssl_show_warn=False,
                request_timeout=30
            )
            if not es.ping():
                print('重新连接失败')
                return
        
        info = es.info()
        print(f"ES连接成功: {info}")
    except Exception as e:
        print(f"ES连接失败: {e}")
        return
    data_in = [
        {
            "_index": index_name,
            "_source": x
        } for x in data
    ]

    success, failed = helpers.bulk(es, data_in)
    # response = es.index(index=index_name, body=data_in)
    print(f'成功写入 {success} 条数据，失败 {failed} 条数据')


def query_data(index_name: str, search_terms: list[str], fields: list[str], permission_type: int = 1,
               permission_set=None, query_file=None, size=10):
    if not es.ping():
        print('无法连接到 Elasticsearch')
        return
    if permission_set is None:
        permission_set = []
    if query_file is None:
        query_file = []

    should_clause = []
    for term in search_terms:
        should_clause.append({
            "multi_match": {"query": term, "fields": fields}
        })
    query = {
        "query": {
            "bool": {
                "should": should_clause
            }
        }
    }
    result = es.search(index=index_name, body=query, size=size)
    dataset = []
    for hit in result['hits']['hits']:
        r = hit['_source']
        r["score"] = hit["_score"]
        # r["permission_set"] =
        if r["file_id"] in query_file:
            dataset.append(r)
        elif permission_type == 1:
            dataset.append(r)
        # elif permission_type == 2 and :
    return dataset


def delete_by_id(index_name, doc_id):
    """通过ID删除文档"""
    try:
        delete_query = {
            "query": {
                "match": {
                    "file_id": doc_id
                }
            }
        }
        response = es.delete_by_query(
            index=index_name,
            body=delete_query
        )
        print(response)
        deleted_count = response['deleted']
        if deleted_count != 0:
            info = f"成功删除{deleted_count}个文档, doc_id: {doc_id}"
            return True, info
        info = f"文档未删除或不存在: {doc_id}"
        return False, info

    except Exception as e:
        print(traceback.format_exc())
        info = f"删除文档失败: {e}"
        return False, info


def get_all_documents(index_name: str = "law1"):
    """获取所有上传的文档列表"""
    global es
    try:
        # 检查连接状态
        if not es.ping():
            print('无法连接到 Elasticsearch，尝试重新连接...')
            # 重新创建连接
            es = Elasticsearch(
                f"{ES_SCHEME}://{ES_USERNAME}:{ES_PASSWORD}@{ES_HOST}:{ES_PORT}",
                verify_certs=False,
                ssl_show_warn=False,
                request_timeout=30
            )
            if not es.ping():
                print('重新连接失败')
                return []
        
        # 使用聚合查询获取所有唯一的file_id和对应的文档信息
        query = {
            "size": 0,
            "aggs": {
                "unique_files": {
                    "terms": {
                        "field": "file_id.keyword",
                        "size": 1000
                    },
                    "aggs": {
                        "doc_info": {
                            "top_hits": {
                                "size": 1,
                                "_source": ["title", "file_id", "permission_set"],
                                "sort": [{"_score": {"order": "desc"}}]
                            }
                        }
                    }
                }
            }
        }
        
        response = es.search(index=index_name, body=query)
        
        documents = []
        for bucket in response['aggregations']['unique_files']['buckets']:
            doc_info = bucket['doc_info']['hits']['hits'][0]['_source']
            documents.append({
                'file_id': doc_info['file_id'],
                'title': doc_info['title'],
                'permission_set': doc_info['permission_set'],
                'doc_count': bucket['doc_count']  # 文档块数量
            })
        
        return documents
    except Exception as e:
        print(f"获取文档列表失败: {e}")
        import traceback
        traceback.print_exc()
        return []


if __name__ == '__main__':
    insert_data("law", [{"title": "中国联通合作方数据安全管理规范", "content": "test", "file_id": "test",
                         "permission_set": [2, 3, 4]}])
    insert_data("law1", [{"title": "中国联通合作方数据安全管理规范", "content": x, "file_id": x[:3],
                          "permission_set": json.dumps([2, 3, 4])} for x in law_data])
    es.indices.refresh(index='law1')
    r = query_data("law1", ["消费者", ], ["title", "content"])
    # r = query_data("law", ["第二十"], ["file_id"])
    print(r)
    # r = delete_by_id("law", "第二十")
    # print(r)
    # es.indices.delete(index='law')
